from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
from dolfinx import mesh, fem
from dolfinx.mesh import locate_entities_boundary, meshtags
from dolfinx.fem.petsc import assemble_matrix, assemble_vector
import ufl
import config as cfg


# ----------------- 基本设置 -----------------
problem_type = "smooth2"  # "smooth1", "smooth2", "singular"
type_ = "SIPG"            # 本代码按 SIPG 正确号法实现
# for NIPG/IIPG: they are suboptimal if the polynomial degree is even
k = 3                     # 多项式阶数

if problem_type == "smooth1":
    x_min, x_max = 0.0, 1.0
    y_min, y_max = 0.0, 1.0
    beta0 = 1.0  # penalization
    
    p_solution = cfg.smooth_solution1
    p_solution_dx = cfg.smooth_solution1_dx
    p_solution_dy = cfg.smooth_solution1_dy
elif problem_type == "smooth2":
    x_min, x_max = 0.0, 1.0
    y_min, y_max = 0.0, 1.0
    beta0 = 3.0  # superpenalization
    p_solution = cfg.smooth_solution2
    p_solution_dx = cfg.smooth_solution2_dx
    p_solution_dy = cfg.smooth_solution2_dy
elif problem_type == "singular":
    x_min, x_max = -1.0, 1.0
    y_min, y_max = -1.0, 1.0
    p_solution = cfg.singular_solution
    p_solution_dx = cfg.singular_solution_dx
    p_solution_dy = cfg.singular_solution_dy
    

if type_ == "NIPG":
    epsilon = 1
    sigma_0 = 1
    sigma_1 = 0
elif type_ == "IIPG":
    epsilon = 0
    sigma_0 = 1
    sigma_1 = 1
elif type_ == "SIPG":
    epsilon = -1   
    sigma_0 = 36
    sigma_1 = 0  # 注意J1项此时为0了

h_list = [1/8, 1/16, 1/32, 1/64, 1/128]
# 误差收敛统计
h1_semi_list, l2_list = [], []
h1_semi_err,  l2_err  = [], []

for h in h_list:
    N = int(1/h)
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [[x_min, y_min], [x_max, y_max]],
        [N, N],
        mesh.CellType.triangle,
    )
    V = fem.functionspace(domain, ("DG", k))
    X = ufl.SpatialCoordinate(domain)
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    
    # -------- 测度与几何量 --------
    n   = ufl.FacetNormal(domain)
    dx  = ufl.dx(domain)
    dS  = ufl.dS(domain=domain)  # 内部面
    # 标记边界
    
    h_cell = ufl.CellDiameter(domain)
    h_avg  = ufl.avg(h_cell)
    

# -------- 收敛阶统计 --------
def _rates(errs, hs):
    out = []
    for i in range(len(errs) - 1):
        out.append(np.log(errs[i] / errs[i+1]) / np.log(hs[i] / hs[i+1]))
    return out

l2_list = _rates(l2_err, h_list)
h1_semi_list = _rates(h1_semi_err, h_list)

if MPI.COMM_WORLD.rank == 0:
    print("\n===== 收敛结果 =====")
    print("h:               ", ["{:.5f}".format(h) for h in h_list])
    print("L2范数误差:       ", ["{:.3e}".format(e) for e in l2_err])
    print("H1半范数误差:     ", ["{:.3e}".format(e) for e in h1_semi_err])
    print("L2收敛阶:         ", ["{:.3f}".format(e) for e in l2_list])
    print("H1半范数收敛阶:   ", ["{:.3f}".format(e) for e in h1_semi_list])